from trectool import read_content, list_content, ParserData
import xml.etree.ElementTree as ET

# Read the title from XML content
def read_title_AP(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('HEAD')
	if title_element is not None:
		return ET.tostring(title_element).replace('<HEAD>','').replace('</HEAD>', '')
	else:
		return None

def read_title_CREFILES(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('TEXT')
	if title_element is not None:
		title = title_element.find('TTL')
	else:
		title = None
	if title is not None:
		return ET.tostring(title).replace('<TTL>','').replace('</TTL>', '')
	else:
		return None

def read_title_CRHFILES(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('TTL')
	if title_element is not None:
		return ET.tostring(title_element).replace('<TTL>','').replace('</TTL>', '')
	else:
		return None

def read_title_FT(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('HEADLINE')
	return ET.tostring(title_element).replace('<HEADLINE>','').replace('</HEADLINE>', '')

def read_title_LATIMES(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('HEADLINE')
	if title_element is not None:
		title = title_element.find('P')
		return ET.tostring(title).replace('<HEADLINE>','').replace('</HEADLINE>', '').replace('<P>', '').replace('</P>', '')
	else:
		return None
	
def read_title_PATENTS(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('TTL')
	if title_element is not None:
		return ET.tostring(title_element).replace('<TTL>','').replace('</TTL>', '')
	else:
		return None

def read_title_SJM(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('DESCRIPT')
	if title_element is not None:
		return ET.tostring(title_element).replace('<DESCRIPT>','').replace('</DESCRIPT>', '')
	else:
		return None

def read_title_WSJ(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('HL')
	if title_element is not None:
		return ET.tostring(title_element).replace('<HL>','').replace('</HL>', '').replace('\n','')
	else:
		return None
	
def read_title_FBIS(xml_content):
	root = ET.fromstring(xml_content)
	title_element = root.find('H3')
	if title_element is not None:
		title = title_element.find('TI')
	else:
		title = None
		
	if title is not None:
		return ET.tostring(title).replace('<TI>','').replace('</TI>', '')
	else:
		return None
			
# Read the text content from XML content
def read_text_content(xml_content):
	root = ET.fromstring(xml_content)
	if root is not None:
		text_element = root.find('TEXT')
	else:
		return None
	if text_element is not None:
		return ET.tostring(text_element).replace('<TEXT>','').replace('</TEXT>', '').replace('<PAR>', '').replace('</PAR>', '').replace('<ABST>', '').replace('</ABST>', '').replace('<BSUM>', '').replace('</BSUM>', '').replace('<DRWD>', '').replace('</DRWD>', '').replace('<DETD>', '').replace('</DETD>', '').replace('<CLMS>', '').replace('</CLMS>', '').replace('<STM>', '').replace('</STM>', '').replace('<NUM>', '').replace('</NUM>', '').replace('<PA1>', '').replace('</PA1>', '').replace('<P>', '').replace('</P>', '').replace('<TTL>','').replace('</TTL>','').replace('<CENTER>','').replace('</CENTER>','').replace('<FLD001>','').replace('</FLD001>','').replace('<FLD002>','').replace('</FLD002>','').replace('<SO>','').replace('</SO>','').replace('<TI>','').replace('</TI>','').replace('<H2>','').replace('</H2>','').replace('<PRE>','').replace('</PRE>','').replace('<PAC>','').replace('</PAC>','').replace('<PARN>','').replace('</PARN>','')
	else:
		return None
	
# Read document and title
def read_document(xml_content, name):
	doc = { 'title' : None, 'content' : None }
	if name == 'AP':
		doc['title'] = read_title_AP(xml_content)
	if name == 'CR_EFILES':
		doc['title'] = read_title_CREFILES(xml_content)
	if name == 'CR_HFILES':
		doc['title'] = read_title_CRHFILES(xml_content)
	if name == 'FBIS':
		doc['title'] = read_title_FBIS(xml_content)
	if name == 'FT':
		doc['title'] = read_title_FT(xml_content)
	if name == 'LATIMES':
		doc['title'] = read_title_LATIMES(xml_content)
	if name == 'PATENTS':
		doc['title'] = read_title_PATENTS(xml_content)
	if name == 'SJM':
		doc['title'] = read_title_SJM(xml_content)
	if name == 'WJM':
		doc['title'] = read_title_WSJ(xml_content)
	doc['content'] = read_text_content(xml_content)
	if doc['title'] is None or doc['content'] is None:
		return None
	else:
		return doc	

'''
name = "AP"
ext = ".txt"
item = name + ".zip"
parser = ParserData(data_root = '..\\testdata')
files = list_content(parser, item, verbose=False)

i = 1
print(len(files))
for a_file in files:
	print(a_file)
	print(len(files))
	xmlcontent = read_content(parser, item, a_file)
	
	#print(read_title_AP(xmlcontent))
	#print(read_title_CREFILES(xmlcontent))
	#print(read_title_CRHFILES(xmlcontent))
	#print(read_title_FT(xmlcontent))
	#print(read_title_LATIMES(xmlcontent))
	#print(read_title_PATENTS(xmlcontent))
	#print(read_title_SJM(xmlcontent))
	#print(read_title_WSJ(xmlcontent))
	#print(read_title_FBIS(xmlcontent))
	
	print(read_docunent(xmlcontent,name))
	
	file_path = "D:\\python\\TextProcessing\\" + name + "\\output" + name + str(i) + ext
	text_file = open("D:\python\TextProcessing\' + name + '\WSJoutput'  + str(i) + '.txt","w")
	text_file = open(file_path,"w")
	text_file.write(read_text_content(xmlcontent))
	text_file.close()
'''
